import { Request, Response } from 'express'
import { PrismaClient } from '@prisma/client'
import WebSocket from 'ws'

const prisma = new PrismaClient()

// 服务器连接池
const serverConnections = new Map<string, WebSocket>()

// 获取服务器列表
export const getServers = async (req: Request, res: Response) => {
    try {
        const servers = await prisma.server.findMany({
            select: {
                id: true,
                name: true,
                host: true,
                status: true,
                gpuInfo: true,
                models: true,
                powerLimit: true,
                priority: true
            }
        })

        // 添加在线状态
        const serversWithStatus = servers.map(server => ({
            ...server,
            isOnline: serverConnections.has(server.id)
        }))

        res.json(serversWithStatus)
    } catch (error: any) {
        res.status(500).json({ error: error.message })
    }
}

// 添加服务器
export const addServer = async (req: Request, res: Response) => {
    try {
        const { name, host, powerLimit, priority } = req.body

        const server = await prisma.server.create({
            data: {
                name,
                host,
                powerLimit,
                priority,
                status: 'offline',
                gpuInfo: {},
                models: []
            }
        })

        res.json(server)
    } catch (error: any) {
        res.status(500).json({ error: error.message })
    }
}

// 更新服务器状态
export const updateServerStatus = async (serverId: string, status: any) => {
    try {
        await prisma.server.update({
            where: { id: serverId },
            data: {
                status: status.status,
                gpuInfo: status.gpuInfo,
                models: status.models
            }
        })
    } catch (error) {
        console.error('更新服务器状态失败:', error)
    }
}

// 处理服务器连接
export const handleServerConnection = (ws: WebSocket, serverId: string) => {
    // 保存连接
    serverConnections.set(serverId, ws)

    // 更新状态为在线
    updateServerStatus(serverId, {
        status: 'online',
        gpuInfo: {},
        models: []
    })

    // 处理消息
    ws.on('message', async (message: string) => {
        try {
            const data = JSON.parse(message)
            if (data.type === 'status') {
                await updateServerStatus(serverId, data.data)
            }
        } catch (error) {
            console.error('处理服务器消息失败:', error)
        }
    })

    // 处理断开连接
    ws.on('close', () => {
        serverConnections.delete(serverId)
        updateServerStatus(serverId, {
            status: 'offline',
            gpuInfo: {},
            models: []
        })
    })
}

// 获取可用服务器
export const getAvailableServer = async (workflowId: string) => {
    try {
        // 获取工作流配置
        const workflow = await prisma.workflow.findUnique({
            where: { id: workflowId },
            select: {
                serverIds: true,
                excludeServerIds: true,
                powerConsumption: true
            }
        })

        if (!workflow) {
            throw new Error('工作流不存在')
        }

        // 构建查询条件
        const where: any = {
            status: 'online'
        }

        // 指定服务器
        if (workflow.serverIds?.length) {
            where.id = { in: workflow.serverIds }
        }

        // 排除服务器
        if (workflow.excludeServerIds?.length) {
            where.id = { notIn: workflow.excludeServerIds }
        }

        // 算力限制
        if (workflow.powerConsumption) {
            where.powerLimit = {
                gte: workflow.powerConsumption
            }
        }

        // 查找服务器
        const servers = await prisma.server.findMany({
            where,
            orderBy: {
                priority: 'desc'
            }
        })

        // 返回第一个可用服务器
        return servers[0] || null

    } catch (error) {
        console.error('获取可用服务器失败:', error)
        return null
    }
} 